import numpy as np
from cyw_devkit.core.visualizer import pad_ones, tranform_pt
from cyw_devkit.core.zmath import rotationMatrixToEulerAngles
from copy import deepcopy


class Header:
    def __init__(self, stamp: int, frame_id: str):
        self.stamp = stamp
        self.frame_id = frame_id


class Data:
    def __init__(self, stamp, frame_id, data):
        self.header = Header(stamp, frame_id)
        self.data = data


class BBox:
    def __init__(self, data):
        if not isinstance(data, dict):
            raise TypeError("Expected a dictionary for BBox data")
        for key, value in data.items():
            setattr(self, key, value)


class BBoxes:
    def __init__(self, labels):
        if not isinstance(labels, list):
            raise TypeError("Expected a list of dictionaries for BBoxes labels")
        self.obj_ids = []
        self.obj_types = []
        self.tensors = []
        for lab in labels:
            self.obj_ids.append(lab["obj_id"])
            self.obj_types.append(lab["obj_type"])
            boxes = [lab['psr']['position']['x'], lab['psr']['position']['y'], lab['psr']['position']['z'],
                     lab['psr']['scale']['x'], lab['psr']['scale']['y'], lab['psr']['scale']['z'],
                     lab['psr']['rotation']['z']]
            self.tensors.append(np.array(boxes))
        self.tensors = np.array(self.tensors)


def transform_bboxes(bboxes, tf):
    new_boxes = deepcopy(bboxes)
    points3 = bboxes.tensors[:, :3]
    points4 = pad_ones(points3)
    new_boxes.tensors[:, :3] = tranform_pt(points4, tf)[:, :3]
    angle = rotationMatrixToEulerAngles(tf[0][:3,:3])[2]
    new_boxes.tensors[:, 6] = new_boxes.tensors[:, 6] + angle
    return new_boxes
